package com.tandashi.springbootstudy01.service.es;

import org.apache.commons.lang3.Validate;
import org.elasticsearch.action.search.MultiSearchRequest;
import org.elasticsearch.action.search.MultiSearchResponse;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.search.aggregations.AggregationBuilder;
import org.elasticsearch.search.aggregations.Aggregations;
import org.elasticsearch.search.aggregations.metrics.NumericMetricsAggregation;
import org.elasticsearch.search.builder.SearchSourceBuilder;

import java.io.IOException;
import java.util.Arrays;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import static com.tandashi.springbootstudy01.convention.constants.ElasticConstant.MAX_FETCH_SIZE;

public abstract class ElasticServiceBase {

    public static SearchRequest buildAggRequest(String index, QueryBuilder query, AggregationBuilder agg) {
        SearchSourceBuilder source = SearchSourceBuilder
                .searchSource()
                .from(0)
                .size(0)
                .query(query)
                .aggregation(agg);

        return buildSearchRequest(index, source);

    }

    private static SearchRequest buildSearchRequest(String index, SearchSourceBuilder source) {
        return new SearchRequest(new String[]{index}, source);
    }

    protected abstract RestHighLevelClient getClient();

    public Map<String, Number> aggSingleValue(MultiSearchRequest request) throws IOException {
        List<Aggregations> aggregationList = aggregations(request);
        return aggregationList
                .stream()
                .map(Aggregations::asList)
                .filter(aggregations -> !aggregations.isEmpty())
                .map(aggregations -> (NumericMetricsAggregation.SingleValue) aggregations.get(0))
                .collect(Collectors.toMap(
                        NumericMetricsAggregation.SingleValue::getName,
                        NumericMetricsAggregation.SingleValue::value,
                        (o, n) -> n,
                        LinkedHashMap::new));
    }

    public List<Aggregations> aggregations(MultiSearchRequest request) throws IOException {
        Validate.isTrue(request.requests().size() <= MAX_FETCH_SIZE,
                "multi request count must not be larger than %s", MAX_FETCH_SIZE);
        for (SearchRequest req : request.requests()) {
            SearchSourceBuilder source = req.source();
            Validate.isTrue(0 == source.from() && 0 == source.size(),
                    "both request from and size must be 0 when aggregation");
            Validate.isTrue(source.aggregations().count() < MAX_FETCH_SIZE,
                    "aggregation count must not be larger than %s", MAX_FETCH_SIZE);
        }

        return Arrays.stream(search(request).getResponses())
                .filter(item -> !item.isFailure())
                .map(MultiSearchResponse.Item::getResponse)
                .map(SearchResponse::getAggregations)
                .collect(Collectors.toList());
    }

    protected MultiSearchResponse search(MultiSearchRequest request) throws IOException {
        return getClient().msearch(request, RequestOptions.DEFAULT);
    }
}
